import os
from PIL import Image
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from os import path
import json
from tqdm import tqdm


class DeepFashionDataset(Dataset):
    def __init__(self, json_path, train=True):
        self.json_path = json_path
        self.train = train
        self.img_dir = path.join(path.dirname(json_path), "images")
        norm = transforms.Normalize([0.831, 0.813, 0.806],[0.229, 0.244, 0.252]) # global mean and std
        # norm = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # imageNet mean and std
        if train:
            self.transform = transforms.Compose(
                [
                    transforms.Resize((256, 256)),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    norm,
                ]
            )
        else:
            self.transform = transforms.Compose(
                [
                    transforms.Resize((256, 256)),
                    transforms.ToTensor(),
                    norm,
                ]
            )
        with open(json_path, "r") as f:
            data = json.load(f)

        self.images_name = list(data.keys())
        self.labels = np.array(list(data.values()))

    def __len__(self):
        return len(self.images_name)

    def __getitem__(self, idx):
        img_name = self.images_name[idx]
        img = Image.open(path.join(self.img_dir, img_name))
        img = self.transform(img)
        return img, self.labels[idx]


def calculate_global_mean_std(img_dir, image_names):
    image_names = np.array(image_names)
    _len = image_names.shape[0]
    means, square = [0 for _ in range(_len)], [0 for _ in range(_len)]
    global_mean, global_std = 0, 0

    for idx, img_name in tqdm(enumerate(image_names), total=_len, desc="Caculating"):
        _img = Image.open(path.join(img_dir, img_name))
        _img = np.asarray(_img, dtype=np.uint32).reshape(-1, 3) / 255  # (H*W, 3), RGB, [0, 1]
        means[idx] = np.mean(_img, axis=0)
        square[idx] = np.square(_img).mean(axis=0)
    global_mean = np.mean(means, axis=0)
    global_std = np.sqrt(np.mean(square, axis=0) - np.square(global_mean))
    return global_mean, global_std
